package org.nerve.webmagic;

import com.alibaba.fastjson.JSON;
import org.apache.commons.lang3.StringUtils;
import org.apache.http.HttpHost;
import org.jsoup.nodes.Document;
import org.nerve.webmagic.bean.Cookie;
import org.nerve.webmagic.bean.KeyValue;
import org.nerve.webmagic.bean.SpiderConfig;
import org.nerve.webmagic.bean.WebPage;
import org.nerve.webmagic.pool.CounterPool;
import org.nerve.webmagic.utils.ImageUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import us.codecraft.webmagic.Page;
import us.codecraft.webmagic.Request;
import us.codecraft.webmagic.Site;
import us.codecraft.webmagic.processor.PageProcessor;

import java.io.IOException;
import java.nio.file.Paths;
import java.util.List;
import java.util.concurrent.atomic.AtomicLong;

/**
 * org.nerve.webmagic.spider.processor
 * Created by zengxm on 2016/4/20 0020.
 */
public class CommonPageProcessor implements PageProcessor {
	private static Logger logger= LoggerFactory.getLogger(CommonPageProcessor.class);

	/**
	 * 成功爬取的网页总数
	 */
	private AtomicLong countPick=new AtomicLong(0L);

	private Site site;
	private final SpiderConfig config;

	private ImageUtil imageUtil;

	public CommonPageProcessor(SpiderConfig config){
		this.config=config;

		imageUtil=new ImageUtil();

		logger.info("page processor created OK:\n{}", JSON.toJSONString(config, true));
	}

	public SpiderConfig getConfig() {
		return config;
	}

	/**
	 * 得到成功获取的网页数
	 * @return
	 */
	public long getPickCount(){
		return countPick.longValue();
	}

	public CommonPageProcessor setScreenDirectory(String screenDirectory) {
		this.config.setScreenDirectory(screenDirectory);
		return this;
	}

	/**
	 * 是否保存截图元数据（base64）
	 * 当保存了截图到文件后，会用到这个字段判断是否清空元数据（清空page中的field）
	 * @param keepScreenshotData
	 * @return
	 */
	public CommonPageProcessor setKeepScreenshotData(boolean keepScreenshotData) {
		this.config.setKeepScreenData(keepScreenshotData);
		return this;
	}
	public CommonPageProcessor setSaveSccreen(boolean saveSccreen) {
		this.config.setSaveScreenToLocal(saveSccreen);
		return this;
	}

	public void process(Page page) {
		//成功爬取的数目+1
//		countPick.incrementAndGet();
		CounterPool.addPick();

		Document doc=page.getHtml().getDocument();
		page.putField(WebPage.TITLE, doc.title());
		page.putField(WebPage.URL, page.getRequest().getUrl());
		page.putField(WebPage.HTML,  doc.toString());
		page.putField(WebPage.DOMAIN, config.getDomain());
		page.putField(WebPage.DRIVER, config.isJbrowser()?"JBrowser":config.isPhantomjs()?"PhantomJS":"");

		//判断是否有截图数据
		String screenData=page.getResultItems().get(WebPage.SCREENSHOT);
		if(screenData!=null){
			try {
				String screenshotPath=Paths.get(config.getScreenDirectory(), System.currentTimeMillis()+".png").toString();
				logger.info("saving screenshot to "+screenshotPath);

				//判断是否需要保存截图到本地
				if(config.isSaveScreenToLocal() && imageUtil.base64StringToImgAndSave(screenData, screenshotPath)){
					page.putField("screenShotPath", screenshotPath);
				}

				if(!config.isKeepScreenData())
					page.putField(WebPage.SCREENSHOT,null);
			} catch (IOException e) {
				logger.error("error on save screenshot:{}", e.getMessage());
			}
		}

		Object level=page.getRequest().getExtra(DeepLimitScheduler.LEVEL);
		int curLevel=level==null?0:((Integer)level);

		List<String> links=page.getHtml().links()
				.regex(config.getLinkRule())
				.all();

		logger.info("process page: title={},links size={}",doc.title(),links.size());

		//构建下一个深度的Request
		links.forEach(l-> page.addTargetRequest(
				new Request(l).putExtra(DeepLimitScheduler.LEVEL, curLevel+1)));

		logger.info("--------------------end process-----------------------");
	}

	public Site getSite() {
		if(site==null){
			site=Site.me()
				.setRetryTimes(3)
				.setUserAgent(config.getUserAgent())
				.setDomain(config.getDomain())
				.setSleepTime(config.getSleepTime());

			if(StringUtils.isNotBlank(config.getCharset()))
				site.setCharset(config.getCharset());

			if(config.getHeaders()!=null){
				for(KeyValue kv:config.getHeaders())
					site.addHeader(kv.name, kv.value);
			}
			else{
				//使用默认的header
				site.addHeader("Connection","keep-alive")
						.addHeader("Accept","text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8");
			}

			if(config.getCookies()!=null)
				for(Cookie cookie:config.getCookies())
					site.addCookie(cookie.domain, cookie.name,cookie.value);

			site.setUseGzip(config.isUseGzip());

			//设置代理
			if(StringUtils.isNotBlank(config.getProxyHost())){
				HttpHost host=config.getProxyPort()>0?
						new HttpHost(config.getProxyHost(),config.getProxyPort())
						:
						new HttpHost(config.getProxyHost());
				site.setHttpProxy(host);
			}
		}

		return site;
	}
}
